retry.py

"""Retry logic with exponential backoff and jitter."""

from __future__ import annotations

import asyncio
import random
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Callable, TypeVar

from .errors import HawkAPIError, RateLimitError

if TYPE_CHECKING:
    from collections.abc import Awaitable

T = TypeVar("T")


@dataclass
class RetryConfig:
    """Configuration for retry behavior."""

    max_retries: int = 3
    initial_backoff: float = 0.5
    max_backoff: float = 30.0
    multiplier: float = 2.0
    retryable_statuses: set[int] = field(default_factory=lambda: {429, 500, 502, 503, 504})


DEFAULT_RETRY_CONFIG = RetryConfig()


def _compute_backoff(attempt: int, config: RetryConfig) -> float:
    """Compute backoff duration with exponential growth and jitter."""
    backoff = config.initial_backoff * (config.multiplier**attempt)
    backoff = min(backoff, config.max_backoff)
    # Add jitter: random value between 0 and backoff
    jitter = random.uniform(0, backoff * 0.5)
    return backoff + jitter


def _is_retryable(error: Exception, config: RetryConfig) -> bool:
    """Determine if an error is retryable based on config."""
    if isinstance(error, HawkAPIError):
        return error.status_code in config.retryable_statuses
    return False


async def with_retry(
    fn: Callable[[], Awaitable[T]],
    config: RetryConfig | None = None,
) -> T:
    """Execute an async function with retry logic.

    Uses exponential backoff with jitter. Respects Retry-After header
    on 429 responses.

    Args:
        fn: Async callable to execute.
        config: Retry configuration. Uses DEFAULT_RETRY_CONFIG if None.

    Returns:
        The result of fn() on success.

    Raises:
        The last exception if all retries are exhausted.
    """
    if config is None:
        config = DEFAULT_RETRY_CONFIG

    last_error: Exception | None = None

    for attempt in range(config.max_retries + 1):
        try:
            return await fn()
        except Exception as e:
            last_error = e

            if attempt >= config.max_retries:
                break

            if not _is_retryable(e, config):
                break

            # Respect Retry-After header for rate limit errors
            if isinstance(e, RateLimitError) and e.retry_after is not None:
                wait = e.retry_after
            else:
                wait = _compute_backoff(attempt, config)

            await asyncio.sleep(wait)

    if last_error is None:
        raise HawkAPIError("Retry logic failed unexpectedly: no error recorded", status_code=0)
    raise last_error


def with_retry_sync(
    fn: Callable[[], T],
    config: RetryConfig | None = None,
) -> T:
    """Execute a sync function with retry logic.

    Uses exponential backoff with jitter. Respects Retry-After header
    on 429 responses.

    Args:
        fn: Callable to execute.
        config: Retry configuration. Uses DEFAULT_RETRY_CONFIG if None.

    Returns:
        The result of fn() on success.

    Raises:
        The last exception if all retries are exhausted.
    """
    if config is None:
        config = DEFAULT_RETRY_CONFIG

    last_error: Exception | None = None

    for attempt in range(config.max_retries + 1):
        try:
            return fn()
        except Exception as e:
            last_error = e

            if attempt >= config.max_retries:
                break

            if not _is_retryable(e, config):
                break

            if isinstance(e, RateLimitError) and e.retry_after is not None:
                wait = e.retry_after
            else:
                wait = _compute_backoff(attempt, config)

            time.sleep(wait)

    if last_error is None:
        raise HawkAPIError("Retry logic failed unexpectedly: no error recorded", status_code=0)
    raise last_error