webber/webber-api/src/shared/retry.py

"""
Retry utilities for handling transient failures.

Provides decorators and helpers for automatic retry with exponential backoff.
"""
import asyncio
import random
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Any, TypeVar

import httpx

from src.shared.logging import get_logger

logger = get_logger(__name__)

T = TypeVar("T")


# Exceptions that should trigger a retry
RETRYABLE_EXCEPTIONS = (
    httpx.TimeoutException,
    httpx.ConnectError,
    httpx.ReadError,
    httpx.WriteError,
    httpx.ConnectTimeout,
    httpx.ReadTimeout,
    httpx.WriteTimeout,
    httpx.PoolTimeout,
    ConnectionError,
    TimeoutError,
    OSError,  # Covers many network-related errors
)


def is_retryable_http_status(status_code: int) -> bool:
    """
    Check if an HTTP status code should trigger a retry.

    Retryable:
    - 429 Too Many Requests (rate limited)
    - 500 Internal Server Error
    - 502 Bad Gateway
    - 503 Service Unavailable
    - 504 Gateway Timeout
    """
    return status_code in (429, 500, 502, 503, 504)


def is_retryable_exception(exc: Exception) -> bool:
    """Check if an exception should trigger a retry."""
    if isinstance(exc, RETRYABLE_EXCEPTIONS):
        return True

    # Check for retryable HTTP status codes
    if isinstance(exc, httpx.HTTPStatusError):
        return is_retryable_http_status(exc.response.status_code)

    return False


def calculate_backoff(
    attempt: int,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter: bool = True,
) -> float:
    """
    Calculate exponential backoff delay with optional jitter.

    Args:
        attempt: Current attempt number (0-indexed)
        base_delay: Base delay in seconds
        max_delay: Maximum delay in seconds
        jitter: Add random jitter to prevent thundering herd

    Returns:
        Delay in seconds
    """
    # Exponential backoff: base_delay * 2^attempt
    delay = min(base_delay * (2 ** attempt), max_delay)

    if jitter:
        # Add up to 25% random jitter
        delay = delay * (0.75 + random.random() * 0.5)

    return delay


def with_retry(
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    retryable_exceptions: tuple[type[Exception], ...] | None = None,
) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
    """
    Decorator for async functions that should retry on transient failures.

    Args:
        max_attempts: Maximum number of attempts (including initial)
        base_delay: Base delay between retries in seconds
        max_delay: Maximum delay between retries in seconds
        retryable_exceptions: Additional exceptions to retry on

    Returns:
        Decorated function with retry logic

    Example:
        @with_retry(max_attempts=3, base_delay=1.0)
        async def fetch_data():
            async with httpx.AsyncClient() as client:
                response = await client.get(url)
                return response.json()
    """
    extra_exceptions = retryable_exceptions or ()

    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> T:
            last_exception: Exception | None = None

            for attempt in range(max_attempts):
                try:
                    return await func(*args, **kwargs)

                except (*RETRYABLE_EXCEPTIONS, *extra_exceptions) as e:
                    last_exception = e
                    should_retry = True

                except httpx.HTTPStatusError as e:
                    last_exception = e
                    should_retry = is_retryable_http_status(e.response.status_code)

                except Exception:
                    # Non-retryable exception, re-raise immediately
                    raise

                if should_retry and attempt < max_attempts - 1:
                    delay = calculate_backoff(attempt, base_delay, max_delay)
                    logger.warning(
                        f"Retry {attempt + 1}/{max_attempts - 1} for {func.__name__} "
                        f"after {delay:.2f}s due to: {last_exception}"
                    )
                    await asyncio.sleep(delay)
                elif not should_retry:
                    # Non-retryable HTTP error
                    raise last_exception  # type: ignore

            # All retries exhausted
            logger.error(
                f"All {max_attempts} attempts failed for {func.__name__}: {last_exception}"
            )
            raise last_exception  # type: ignore

        return wrapper
    return decorator


async def retry_async(
    func: Callable[..., Awaitable[T]],
    *args: Any,
    max_attempts: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    **kwargs: Any,
) -> T:
    """
    Retry an async function with exponential backoff.

    Alternative to decorator when you need per-call control.

    Args:
        func: Async function to call
        *args: Positional arguments for func
        max_attempts: Maximum number of attempts
        base_delay: Base delay between retries
        max_delay: Maximum delay between retries
        **kwargs: Keyword arguments for func

    Returns:
        Result of func

    Raises:
        Last exception if all retries fail

    Example:
        result = await retry_async(
            fetch_data,
            url,
            max_attempts=5,
            timeout=30,
        )
    """
    last_exception: Exception | None = None

    for attempt in range(max_attempts):
        try:
            return await func(*args, **kwargs)

        except Exception as e:
            last_exception = e

            if not is_retryable_exception(e):
                raise

            if attempt < max_attempts - 1:
                delay = calculate_backoff(attempt, base_delay, max_delay)
                logger.warning(
                    f"Retry {attempt + 1}/{max_attempts - 1} "
                    f"after {delay:.2f}s due to: {e}"
                )
                await asyncio.sleep(delay)

    raise last_exception  # type: ignore