118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
"""
|
|
retry.py
|
|
--------
|
|
Exponential backoff retry helpers for API calls that may hit rate limits
|
|
or transient server errors (429, 502, 503, 504).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import functools
|
|
import logging
|
|
import time
|
|
from typing import Callable, TypeVar
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
|
|
# HTTP status codes that are safe to retry
|
|
_RETRYABLE_STATUS = {429, 500, 502, 503, 504}
|
|
|
|
|
|
class RetryableHTTPError(Exception):
|
|
"""Raised for HTTP status codes that warrant a retry (429, 500, 502, 503, 504)."""
|
|
|
|
|
|
def raise_for_retryable_status(resp) -> None:
|
|
"""
|
|
Raise RetryableHTTPError for retryable status codes; call raise_for_status() for
|
|
all others. Use this instead of resp.raise_for_status() in functions decorated with
|
|
@retry_with_backoff(retryable_exceptions=(RetryableHTTPError,)).
|
|
"""
|
|
if resp.status_code in _RETRYABLE_STATUS:
|
|
raise RetryableHTTPError(f"HTTP {resp.status_code} from {resp.url} — will retry")
|
|
resp.raise_for_status()
|
|
|
|
|
|
def retry_with_backoff(
|
|
max_retries: int = 3,
|
|
base_delay: float = 1.0,
|
|
max_delay: float = 30.0,
|
|
retryable_exceptions: tuple = (Exception,),
|
|
):
|
|
"""
|
|
Decorator for sync functions. Retries on exceptions with exponential backoff.
|
|
|
|
Usage:
|
|
@retry_with_backoff(max_retries=3, base_delay=1.0)
|
|
def my_api_call():
|
|
...
|
|
"""
|
|
def decorator(fn: Callable) -> Callable:
|
|
@functools.wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
last_exc: Exception | None = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
except retryable_exceptions as exc:
|
|
last_exc = exc
|
|
if attempt == max_retries:
|
|
break
|
|
delay = min(base_delay * (2 ** attempt), max_delay)
|
|
logger.warning(
|
|
"Retry %d/%d for %s after %.1fs — %s",
|
|
attempt + 1, max_retries, fn.__name__, delay, exc,
|
|
)
|
|
time.sleep(delay)
|
|
raise last_exc
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
def async_retry_with_backoff(
|
|
max_retries: int = 3,
|
|
base_delay: float = 1.0,
|
|
max_delay: float = 30.0,
|
|
retryable_exceptions: tuple = (Exception,),
|
|
):
|
|
"""
|
|
Decorator for async functions. Retries on exceptions with exponential backoff.
|
|
|
|
Usage:
|
|
@async_retry_with_backoff(max_retries=3, base_delay=1.0)
|
|
async def my_api_call():
|
|
...
|
|
"""
|
|
def decorator(fn: Callable) -> Callable:
|
|
@functools.wraps(fn)
|
|
async def wrapper(*args, **kwargs):
|
|
last_exc: Exception | None = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return await fn(*args, **kwargs)
|
|
except retryable_exceptions as exc:
|
|
last_exc = exc
|
|
if attempt == max_retries:
|
|
break
|
|
delay = min(base_delay * (2 ** attempt), max_delay)
|
|
logger.warning(
|
|
"Async retry %d/%d for %s after %.1fs — %s",
|
|
attempt + 1, max_retries, fn.__name__, delay, exc,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
raise last_exc
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
class RateLimitError(Exception):
|
|
"""Raised when an API returns HTTP 429 Too Many Requests."""
|
|
|
|
|
|
def check_response_retryable(status_code: int) -> bool:
|
|
"""Return True if the HTTP status code warrants a retry."""
|
|
return status_code in _RETRYABLE_STATUS
|