chore: consolidate duplicate backoff logic into Backoff class (#2660)

Author: haakonvt

cognite/client/_http_client.py

@@ -3,7 +3,6 @@
 import asyncio
 import functools
 import logging
-import random
 from collections.abc import (
     AsyncIterable,
     AsyncIterator,
@@ -28,6 +27,7 @@
     CogniteRequestError,
 )
 from cognite.client.response import CogniteHTTPResponse
+from cognite.client.utils._retry import Backoff
 
 logger = logging.getLogger(__name__)
 
@@ -132,24 +132,18 @@ def __init__(self, url: str, config: AsyncHTTPClientWithRetryConfig) -> None:
         self.status = self.read = self.connect = 0
         self.last_failed_reason = ""
         self.total_time_in_backoff = 0.0
+        self._backoff = Backoff(
+            multiplier=config.backoff_factor,
+            max_wait=float(config.max_backoff_seconds),
+            min_wait=0.1,
+        )
 
     @property
     def total(self) -> int:
         return self.status + self.read + self.connect
 
-    def get_backoff_time(self) -> float:
-        """
-        Computes the randomized delay (backoff time) before the next retry attempt.
-
-        This method implements the **Exponential Backoff with Full Jitter** strategy,
-        with one addition, a small "initial floor" to avoid immediate retries.
-        """
-        base = self.config.backoff_factor * 2**self.total
-        cap = min(base, self.config.max_backoff_seconds)
-        return random.uniform(0.1, cap)
-
     async def back_off(self) -> None:
-        backoff_time = self.get_backoff_time()
+        backoff_time = next(self._backoff)
         self.total_time_in_backoff += backoff_time
         # We put logging here, since this is always called before retrying
         logger.debug(

cognite/client/utils/_retry.py

@@ -9,20 +9,24 @@ class Backoff(Iterator[float]):
     described in this post: https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
 
     Args:
-        multiplier (float): No description.
-        max_wait (float): No description.
-        base (int): No description."""
+        multiplier (float): Scales the exponential curve; the cap after n attempts is ``multiplier * base**n``.
+        max_wait (float): Hard ceiling on the sampled wait in seconds.
+        base (int): Exponential base; 2 gives classic doubling behaviour.
+        min_wait (float): Lower bound of the uniform sample, useful when an immediate retry is undesirable."""
 
-    def __init__(self, multiplier: float = 0.5, max_wait: float = 60.0, base: int = 2) -> None:
+    def __init__(self, multiplier: float = 0.5, max_wait: float = 60.0, base: int = 2, min_wait: float = 0.0) -> None:
         self._multiplier = multiplier
         self._max_wait = max_wait
         self._base = base
+        self._min_wait = min_wait
         self._past_attempts = 0
 
     def __next__(self) -> float:
         # 100 is an arbitrary limit at which point most sensible parameters are likely to
         # be capped by max anyway.
-        wait = uniform(0, min(self._multiplier * (self._base ** min(100, self._past_attempts)), self._max_wait))
+        base = self._multiplier * (self._base ** min(100, self._past_attempts))
+        cap = max(self._min_wait, min(base, self._max_wait))
+        wait = uniform(self._min_wait, cap)
         self._past_attempts += 1
         return wait
 

tests/tests_unit/test_http_client.py

@@ -105,12 +105,6 @@ def test_correct_retry_of_status(self, default_config: AsyncHTTPClientWithRetryC
         rt.config.status_codes_to_retry.add(500)
         assert rt.should_retry_status_code(make_http_status_error(500)) is True
 
-    def test_get_backoff_time(self, default_config: AsyncHTTPClientWithRetryConfig) -> None:
-        rt = RetryTracker(URL, default_config)
-        for i in range(10):
-            rt.read += 1
-            assert 0 <= rt.get_backoff_time() <= default_config.max_backoff_seconds
-
     def test_is_auto_retryable(self, default_config: AsyncHTTPClientWithRetryConfig) -> None:
         default_config._max_retries_status = 1
         rt = RetryTracker(URL, default_config)