walrasquant-lib 0.4.20__py3-none-any.whl

walrasquant/base/retry.py

@@ -0,0 +1,220 @@
+import asyncio
+import nexuslog as logging
+
+from typing import Callable, Awaitable, Generic, TypeVar
+from random import randint
+
+T = TypeVar("T")
+
+
+def get_exponential_backoff(
+    num_attempts: int,
+    delay_initial_ms: int = 500,
+    delay_max_ms: int = 2_000,
+    backoff_factor: int = 2,
+    jitter: bool = True,
+) -> int:
+    """
+    Compute the backoff using exponential backoff and jitter.
+
+    Parameters
+    ----------
+    num_attempts : int, default 1
+        The number of attempts that have already been made.
+    delay_initial_ms : int, default 500
+        The time to sleep in the first attempt.
+    delay_max_ms : int, default 2_000
+        The maximum delay.
+    backoff_factor : int, default 2
+        The exponential backoff factor for delays.
+    jitter : bool, default True
+        Whether or not to apply jitter.
+
+    Notes
+    -----
+    https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
+
+    Returns
+    -------
+    int
+        Delay in milliseconds.
+
+    """
+    delay = min(delay_max_ms, delay_initial_ms * backoff_factor ** (num_attempts - 1))
+
+    if jitter:
+        return randint(delay_initial_ms, delay)  # noqa: S311
+
+    return delay
+
+
+class RetryManager(Generic[T]):
+    """
+    Provides retry state management for an HTTP request.
+
+    This class is generic over `T`, where `T` is the return type of the
+    function passed to the `run` method.
+
+    Parameters
+    ----------
+    max_retries : int
+        The maximum number of retries before failure.
+    delay_initial_ms : int
+        The initial delay (milliseconds) for retries.
+    delay_max_ms : int
+        The maximum delay (milliseconds) for exponential backoff.
+    backoff_factor : int
+        The exponential backoff factor for retry delays.
+    exc_types : tuple[Type[BaseException], ...]
+        The exception types to handle for retries.
+    retry_check : Callable[[BaseException], None], optional
+        A function that performs additional checks on the exception.
+        If the function returns `False`, a retry will not be attempted.
+
+    """
+
+    def __init__(
+        self,
+        max_retries: int,
+        delay_initial_ms: int,
+        delay_max_ms: int,
+        backoff_factor: int,
+        exc_types: tuple[type[BaseException], ...],
+        retry_check: Callable[[BaseException], bool] | None = None,
+    ) -> None:
+        self.max_retries = max_retries
+        self.delay_initial_ms = delay_initial_ms
+        self.delay_max_ms = delay_max_ms
+        self.backoff_factor = backoff_factor
+        self.exc_types = exc_types
+        self.retry_check = retry_check
+        self.cancel_event = asyncio.Event()
+        self._log = logging.getLogger(name=type(self).__name__)
+
+        self.name: str | None = None
+        self.details: list[object] | None = None
+        self.details_str: str | None = None
+        self.result: bool = False
+        self.message: str | None = None
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__}(name='{self.name}', details={self.details}) at {hex(id(self))}>"
+
+    async def run(
+        self,
+        name: str,
+        func: Callable[..., Awaitable[T]],
+        details: list[object] | None = None,
+        jitter: bool = True,
+        *args,
+        **kwargs,
+    ) -> T | None:
+        """
+        Execute the given `func` with retry management.
+
+        If an exception in `self.exc_types` is raised, a warning is logged, and the function is
+        retried after a delay until the maximum retries are reached, at which point an error is logged.
+
+        Parameters
+        ----------
+        name : str
+            The name of the operation to run.
+        details : list[object], optional
+            The operation details such as identifiers.
+        func : Callable[..., Awaitable[T]]
+            The function to execute.
+        jitter : bool, default True
+            Whether to apply jitter to the retry delay.
+        args : Any
+            Positional arguments to pass to the function `func`.
+        kwargs : Any
+            Keyword arguments to pass to the function `func`.
+
+        Returns
+        -------
+        T | None
+            The result of the executed function, or ``None`` if the retries fail.
+
+        """
+        retries = 0
+        self.name = name
+        self.details = details
+
+        try:
+            while True:
+                if self.cancel_event.is_set():
+                    self._cancel()
+                    return None
+
+                try:
+                    response = await func(*args, **kwargs)
+                    self.result = True
+
+                    return response  # Successful request
+                except self.exc_types as e:
+                    self._log.warning(repr(e))
+                    if (
+                        (self.retry_check and not self.retry_check(e))
+                        or not self.max_retries
+                        or retries >= self.max_retries
+                    ):
+                        self._log_error()
+                        self.result = False
+                        self.message = str(e)
+                        raise e
+
+                    retries += 1
+                    retry_delay_ms = get_exponential_backoff(
+                        delay_initial_ms=self.delay_initial_ms,
+                        delay_max_ms=self.delay_max_ms,
+                        backoff_factor=self.backoff_factor,
+                        num_attempts=retries,
+                        jitter=jitter,
+                    )
+                    self._log_retry(retries, retry_delay_ms=retry_delay_ms)
+                    await asyncio.sleep(retry_delay_ms / 1000)
+        except asyncio.CancelledError:
+            self._cancel()
+            return None
+
+    def cancel(self) -> None:
+        """
+        Cancel the retry operation.
+        """
+        self._log.debug(f"Canceling {self!r}")
+        self.cancel_event.set()
+
+    def clear(self) -> None:
+        """
+        Clear all state from this retry manager.
+        """
+        self.name = None
+        self.details = None
+        self.details_str = None
+        self.result = False
+        self.message = None
+
+    def _cancel(self) -> None:
+        self._log.warning(f"Canceled retry for '{self.name}'")
+        self.result = False
+        self.message = "Canceled retry"
+
+    def _log_retry(self, retry: int, retry_delay_ms: int) -> None:
+        self._log.warning(
+            f"Retrying {retry}/{self.max_retries} for '{self.name}' "
+            f"in {retry_delay_ms / 1000}s{self._details_str()}",
+        )
+
+    def _log_error(self) -> None:
+        self._log.error(
+            f"Failed on '{self.name}'{self._details_str()}",
+        )
+
+    def _details_str(self) -> str:
+        if not self.details:
+            return ""
+
+        if not self.details_str:
+            self.details_str = ": " + ", ".join([repr(x) for x in self.details])
+
+        return self.details_str