retries.py

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

import asyncio
import random
import time
from typing import List

import httpx


class BackoffStrategy:
    initial_interval: int
    max_interval: int
    exponent: float
    max_elapsed_time: int
    min_attempts: int
    absolute_max_elapsed_time_ms: int | None

    def __init__(
        self,
        initial_interval: int,
        max_interval: int,
        exponent: float,
        max_elapsed_time: int,
        min_attempts: int = 0,
        absolute_max_elapsed_time_ms: int | None = None,
    ):
        # min_attempts counts retries (not the initial attempt).
        # absolute_max_elapsed_time_ms caps when a new retry can start;
        # it does not interrupt in-flight func() calls.
        if min_attempts < 0:
            raise ValueError(
                f"min_attempts must be >= 0, got {min_attempts}"
            )
        if absolute_max_elapsed_time_ms is not None:
            if absolute_max_elapsed_time_ms <= 0:
                raise ValueError(
                    "absolute_max_elapsed_time_ms must be > 0, "
                    f"got {absolute_max_elapsed_time_ms}"
                )
            if absolute_max_elapsed_time_ms < max_elapsed_time:
                raise ValueError(
                    "absolute_max_elapsed_time_ms must be >= max_elapsed_time "
                    f"({max_elapsed_time}), got {absolute_max_elapsed_time_ms}"
                )
        self.initial_interval = initial_interval
        self.max_interval = max_interval
        self.exponent = exponent
        self.max_elapsed_time = max_elapsed_time
        self.min_attempts = min_attempts
        self.absolute_max_elapsed_time_ms = absolute_max_elapsed_time_ms


class RetryConfig:
    strategy: str
    backoff: BackoffStrategy
    retry_connection_errors: bool

    def __init__(
        self, strategy: str, backoff: BackoffStrategy, retry_connection_errors: bool
    ):
        self.strategy = strategy
        self.backoff = backoff
        self.retry_connection_errors = retry_connection_errors


class Retries:
    config: RetryConfig
    status_codes: List[str]

    def __init__(self, config: RetryConfig, status_codes: List[str]):
        self.config = config
        self.status_codes = status_codes


class TemporaryError(Exception):
    response: httpx.Response

    def __init__(self, response: httpx.Response):
        self.response = response


class PermanentError(Exception):
    inner: Exception

    def __init__(self, inner: Exception):
        self.inner = inner


def _declares_non_retryable(response: httpx.Response) -> bool:
    # Inspect buffered error bodies without consuming unread response streams.
    try:
        body = response.json()
    except (ValueError, RecursionError, httpx.ResponseNotRead):
        return False
    return isinstance(body, dict) and body.get("retryable") is False


def _is_retry_status(response: httpx.Response, status_codes: List[str]) -> bool:
    for code in status_codes:
        if "X" in code.upper():
            code_range = int(code[0])
            status_major = response.status_code / 100
            if code_range <= status_major < code_range + 1:
                return True
        elif response.status_code == int(code):
            return True
    return False


def retry(func, retries: Retries):
    if retries.config.strategy == "backoff":

        def do_request() -> httpx.Response:
            res: httpx.Response
            try:
                res = func()

                if _is_retry_status(res, retries.status_codes) and not _declares_non_retryable(res):
                    raise TemporaryError(res)
            except httpx.TransportError as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except TemporaryError:
                raise
            except Exception as exception:
                raise PermanentError(exception) from exception

            return res

        return retry_with_backoff(
            do_request,
            retries.config.backoff.initial_interval,
            retries.config.backoff.max_interval,
            retries.config.backoff.exponent,
            retries.config.backoff.max_elapsed_time,
            retries.config.backoff.min_attempts,
            retries.config.backoff.absolute_max_elapsed_time_ms,
        )

    return func()


async def retry_async(func, retries: Retries):
    if retries.config.strategy == "backoff":

        async def do_request() -> httpx.Response:
            res: httpx.Response
            try:
                res = await func()

                if _is_retry_status(res, retries.status_codes) and not _declares_non_retryable(res):
                    raise TemporaryError(res)
            except httpx.TransportError as exception:
                if retries.config.retry_connection_errors:
                    raise

                raise PermanentError(exception) from exception
            except TemporaryError:
                raise
            except Exception as exception:
                raise PermanentError(exception) from exception

            return res

        return await retry_with_backoff_async(
            do_request,
            retries.config.backoff.initial_interval,
            retries.config.backoff.max_interval,
            retries.config.backoff.exponent,
            retries.config.backoff.max_elapsed_time,
            retries.config.backoff.min_attempts,
            retries.config.backoff.absolute_max_elapsed_time_ms,
        )

    return await func()


def _cap_hit_after_attempt(
    elapsed_ms,
    retries,
    min_attempts,
    max_elapsed_time,
    absolute_max_elapsed_time_ms,
):
    soft_cap_hit = retries >= min_attempts and elapsed_ms > max_elapsed_time
    hard_cap_hit = (
        absolute_max_elapsed_time_ms is not None
        and elapsed_ms > absolute_max_elapsed_time_ms
    )
    return soft_cap_hit or hard_cap_hit


def _raise_or_return_after_cap(exception, elapsed_ms, retries, reason):
    if isinstance(exception, TemporaryError):
        return exception.response
    elapsed_seconds = elapsed_ms / 1000
    raise type(exception)(
        f"{type(exception).__name__} after {retries + 1} attempts "
        f"over {elapsed_seconds:.1f}s{reason}: {exception}"
    ) from exception


def retry_with_backoff(
    func,
    initial_interval=500,
    max_interval=60000,
    exponent=1.5,
    max_elapsed_time=3600000,
    min_attempts=0,
    absolute_max_elapsed_time_ms=None,
):
    start = round(time.time() * 1000)
    retries = 0

    while True:
        try:
            return func()
        except PermanentError as exception:
            raise exception.inner
        except Exception as exception:  # pylint: disable=broad-exception-caught
            now = round(time.time() * 1000)
            elapsed = now - start

            if _cap_hit_after_attempt(
                elapsed,
                retries,
                min_attempts,
                max_elapsed_time,
                absolute_max_elapsed_time_ms,
            ):
                result = _raise_or_return_after_cap(
                    exception, elapsed, retries, ""
                )
                return result

            sleep_seconds = (
                (initial_interval / 1000) * exponent**retries
                + random.uniform(0, 1)
            )
            sleep_seconds = min(sleep_seconds, max_interval / 1000)

            if absolute_max_elapsed_time_ms is not None:
                projected_elapsed = elapsed + int(sleep_seconds * 1000)
                if projected_elapsed >= absolute_max_elapsed_time_ms:
                    result = _raise_or_return_after_cap(
                        exception,
                        elapsed,
                        retries,
                        " (hard cap would be exceeded during backoff)",
                    )
                    return result

            time.sleep(sleep_seconds)

            if absolute_max_elapsed_time_ms is not None:
                actual_elapsed = round(time.time() * 1000) - start
                if actual_elapsed >= absolute_max_elapsed_time_ms:
                    result = _raise_or_return_after_cap(
                        exception,
                        actual_elapsed,
                        retries,
                        " (hard cap reached during backoff sleep)",
                    )
                    return result

            retries += 1


async def retry_with_backoff_async(
    func,
    initial_interval=500,
    max_interval=60000,
    exponent=1.5,
    max_elapsed_time=3600000,
    min_attempts=0,
    absolute_max_elapsed_time_ms=None,
):
    start = round(time.time() * 1000)
    retries = 0

    while True:
        try:
            return await func()
        except PermanentError as exception:
            raise exception.inner
        except Exception as exception:  # pylint: disable=broad-exception-caught
            now = round(time.time() * 1000)
            elapsed = now - start

            if _cap_hit_after_attempt(
                elapsed,
                retries,
                min_attempts,
                max_elapsed_time,
                absolute_max_elapsed_time_ms,
            ):
                result = _raise_or_return_after_cap(
                    exception, elapsed, retries, ""
                )
                return result

            sleep_seconds = (
                (initial_interval / 1000) * exponent**retries
                + random.uniform(0, 1)
            )
            sleep_seconds = min(sleep_seconds, max_interval / 1000)

            if absolute_max_elapsed_time_ms is not None:
                projected_elapsed = elapsed + int(sleep_seconds * 1000)
                if projected_elapsed >= absolute_max_elapsed_time_ms:
                    result = _raise_or_return_after_cap(
                        exception,
                        elapsed,
                        retries,
                        " (hard cap would be exceeded during backoff)",
                    )
                    return result

            await asyncio.sleep(sleep_seconds)

            if absolute_max_elapsed_time_ms is not None:
                actual_elapsed = round(time.time() * 1000) - start
                if actual_elapsed >= absolute_max_elapsed_time_ms:
                    result = _raise_or_return_after_cap(
                        exception,
                        actual_elapsed,
                        retries,
                        " (hard cap reached during backoff sleep)",
                    )
                    return result

            retries += 1