added TrackedBackoffGenerator

Author: daniel-sanche

google/cloud/bigtable/data/_helpers.py

@@ -23,6 +23,7 @@
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 
 from google.api_core import exceptions as core_exceptions
+from google.api_core.retry import exponential_sleep_generator
 from google.api_core.retry import RetryFailureReason
 from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
 
@@ -248,3 +249,32 @@ def _get_retryable_errors(
         call_codes = table.default_mutate_rows_retryable_errors
 
     return [_get_error_type(e) for e in call_codes]
+
+
+class TrackedBackoffGenerator:
+    """
+    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
+        )
+
+    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]

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

@@ -32,7 +32,7 @@
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
-    from google.cloud.bigtable.data._helpers import BackoffGenerator
+    from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 
 
 LOGGER = logging.getLogger(__name__)
@@ -144,7 +144,7 @@ class ActiveOperationMetric:
 
     op_type: OperationType
     uuid: str = str(uuid.uuid4())
-    backoff_generator: BackoffGenerator | None = None
+    backoff_generator: TrackedBackoffGenerator | None = None
     # keep monotonic timestamps for active operations
     start_time_ns: int = field(default_factory=time.monotonic_ns)
     active_attempt: ActiveAttemptMetric | None = None

tests/unit/data/test__helpers.py

@@ -266,3 +266,53 @@ def test_get_retryable_errors(self, input_codes, input_table, expected):
             setattr(fake_table, f"{key}_retryable_errors", input_table[key])
         result = _helpers._get_retryable_errors(input_codes, fake_table)
         assert result == expected
+
+
+class TestTrackedBackoffGenerator:
+
+    def test_tracked_backoff_generator_history(self):
+        """
+        Should be able to retrieve historical results from backoff generator
+        """
+        generator = _helpers.TrackedBackoffGenerator(initial=0, multiplier=2, maximum=10)
+        got_list = [next(generator) for _ in range(20)]
+
+        # check all values are correct
+        for i in range(19, 0, -1):
+            assert generator.get_attempt_backoff(i) == got_list[i]
+        # check a random value out of order
+        assert generator.get_attempt_backoff(5) == got_list[5]
+
+    @mock.patch("random.uniform", side_effect=lambda a, b: b)
+    def test_tracked_backoff_generator_defaults(self, mock_uniform):
+        """
+        Should generate values with default parameters
+
+        initial=0.01, multiplier=2, maximum=60
+        """
+        generator = _helpers.TrackedBackoffGenerator()
+        expected_values = [0.01, 0.02, 0.04, 0.08, 0.16]
+        for expected in expected_values:
+            assert next(generator) == pytest.approx(expected)
+
+    @mock.patch("random.uniform", side_effect=lambda a, b: b)
+    def test_tracked_backoff_generator_with_maximum(self, mock_uniform):
+        """
+        Should cap the backoff at the maximum value
+        """
+        generator = _helpers.TrackedBackoffGenerator(initial=1, multiplier=2, maximum=5)
+        expected_values = [1, 2, 4, 5, 5, 5]
+        for expected in expected_values:
+            assert next(generator) == expected
+
+    def test_get_attempt_backoff_out_of_bounds(self):
+        """
+        get_attempt_backoff should raise IndexError for out of bounds index
+        """
+        generator = _helpers.TrackedBackoffGenerator()
+        next(generator)
+        next(generator)
+        with pytest.raises(IndexError):
+            generator.get_attempt_backoff(2)
+        with pytest.raises(IndexError):
+            generator.get_attempt_backoff(-3)