230 lines
7.0 KiB
Python
230 lines
7.0 KiB
Python
"""
|
|
Reverse retry utilities.
|
|
"""
|
|
|
|
import asyncio
|
|
import random
|
|
from typing import Callable, Any, Optional
|
|
|
|
from app.core.logger import logger
|
|
from app.core.config import get_config
|
|
from app.core.exceptions import UpstreamException
|
|
|
|
|
|
class RetryContext:
|
|
"""Retry context."""
|
|
|
|
def __init__(self):
|
|
self.attempt = 0
|
|
self.max_retry = int(get_config("retry.max_retry"))
|
|
self.retry_codes = get_config("retry.retry_status_codes")
|
|
self.last_error = None
|
|
self.last_status = None
|
|
self.total_delay = 0.0
|
|
self.retry_budget = float(get_config("retry.retry_budget"))
|
|
|
|
# Backoff parameters
|
|
self.backoff_base = float(get_config("retry.retry_backoff_base"))
|
|
self.backoff_factor = float(get_config("retry.retry_backoff_factor"))
|
|
self.backoff_max = float(get_config("retry.retry_backoff_max"))
|
|
|
|
# Decorrelated jitter state
|
|
self._last_delay = self.backoff_base
|
|
|
|
def should_retry(self, status_code: int) -> bool:
|
|
"""Check if should retry."""
|
|
if self.attempt >= self.max_retry:
|
|
return False
|
|
if status_code not in self.retry_codes:
|
|
return False
|
|
if self.total_delay >= self.retry_budget:
|
|
return False
|
|
return True
|
|
|
|
def record_error(self, status_code: int, error: Exception):
|
|
"""Record error information."""
|
|
self.last_status = status_code
|
|
self.last_error = error
|
|
self.attempt += 1
|
|
|
|
def calculate_delay(self, status_code: int, retry_after: Optional[float] = None) -> float:
|
|
"""
|
|
Calculate backoff delay time.
|
|
|
|
Args:
|
|
status_code: HTTP status code
|
|
retry_after: Retry-After header value (seconds)
|
|
|
|
Returns:
|
|
Delay time (seconds)
|
|
"""
|
|
# Use Retry-After if available
|
|
if retry_after is not None and retry_after > 0:
|
|
delay = min(retry_after, self.backoff_max)
|
|
self._last_delay = delay
|
|
return delay
|
|
|
|
# Use decorrelated jitter for 429
|
|
if status_code == 429:
|
|
# decorrelated jitter: delay = random(base, last_delay * 3)
|
|
delay = random.uniform(self.backoff_base, self._last_delay * 3)
|
|
delay = min(delay, self.backoff_max)
|
|
self._last_delay = delay
|
|
return delay
|
|
|
|
# Use exponential backoff + full jitter for other status codes
|
|
exp_delay = self.backoff_base * (self.backoff_factor**self.attempt)
|
|
delay = random.uniform(0, min(exp_delay, self.backoff_max))
|
|
return delay
|
|
|
|
def record_delay(self, delay: float):
|
|
"""Record delay time."""
|
|
self.total_delay += delay
|
|
|
|
|
|
def extract_retry_after(error: Exception) -> Optional[float]:
|
|
"""
|
|
Extract Retry-After value from exception.
|
|
|
|
Args:
|
|
error: Exception object
|
|
|
|
Returns:
|
|
Retry-After value (seconds), or None
|
|
"""
|
|
if not isinstance(error, UpstreamException):
|
|
return None
|
|
|
|
details = error.details or {}
|
|
|
|
# Try to get Retry-After from details
|
|
retry_after = details.get("retry_after")
|
|
if retry_after is not None:
|
|
try:
|
|
return float(retry_after)
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
# Try to get Retry-After from headers
|
|
headers = details.get("headers", {})
|
|
if isinstance(headers, dict):
|
|
retry_after = headers.get("Retry-After") or headers.get("retry-after")
|
|
if retry_after is not None:
|
|
try:
|
|
return float(retry_after)
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
async def retry_on_status(
|
|
func: Callable,
|
|
*args,
|
|
extract_status: Callable[[Exception], Optional[int]] = None,
|
|
on_retry: Callable[[int, int, Exception, float], None] = None,
|
|
**kwargs,
|
|
) -> Any:
|
|
"""
|
|
Generic retry function.
|
|
|
|
Args:
|
|
func: Retry function
|
|
*args: Function arguments
|
|
extract_status: Function to extract status code from exception
|
|
on_retry: Callback function for retry (attempt, status_code, error, delay)
|
|
**kwargs: Function keyword arguments
|
|
|
|
Returns:
|
|
Function execution result
|
|
|
|
Raises:
|
|
Last failed exception
|
|
"""
|
|
ctx = RetryContext()
|
|
|
|
# Status code extractor
|
|
if extract_status is None:
|
|
|
|
def extract_status(e: Exception) -> Optional[int]:
|
|
if isinstance(e, UpstreamException):
|
|
# Try to get status code from details, fallback to status_code attribute
|
|
if e.details and "status" in e.details:
|
|
return e.details["status"]
|
|
return getattr(e, "status_code", None)
|
|
return None
|
|
|
|
while ctx.attempt <= ctx.max_retry:
|
|
try:
|
|
result = await func(*args, **kwargs)
|
|
|
|
# Record log
|
|
if ctx.attempt > 0:
|
|
logger.info(
|
|
f"Retry succeeded after {ctx.attempt} attempts, "
|
|
f"total delay: {ctx.total_delay:.2f}s"
|
|
)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
# Extract status code
|
|
status_code = extract_status(e)
|
|
|
|
if status_code is None:
|
|
# Error cannot be identified as retryable
|
|
logger.error(f"Non-retryable error: {e}")
|
|
raise
|
|
|
|
# Record error
|
|
ctx.record_error(status_code, e)
|
|
|
|
# Check if should retry
|
|
if ctx.should_retry(status_code):
|
|
# Extract Retry-After
|
|
retry_after = extract_retry_after(e)
|
|
|
|
# Calculate delay
|
|
delay = ctx.calculate_delay(status_code, retry_after)
|
|
|
|
# Check if exceeds budget
|
|
if ctx.total_delay + delay > ctx.retry_budget:
|
|
logger.warning(
|
|
f"Retry budget exhausted: {ctx.total_delay:.2f}s + {delay:.2f}s > {ctx.retry_budget}s"
|
|
)
|
|
raise
|
|
|
|
ctx.record_delay(delay)
|
|
|
|
logger.warning(
|
|
f"Retry {ctx.attempt}/{ctx.max_retry} for status {status_code}, "
|
|
f"waiting {delay:.2f}s (total: {ctx.total_delay:.2f}s)"
|
|
+ (f", Retry-After: {retry_after}s" if retry_after else "")
|
|
)
|
|
|
|
# Callback
|
|
if on_retry:
|
|
on_retry(ctx.attempt, status_code, e, delay)
|
|
|
|
await asyncio.sleep(delay)
|
|
continue
|
|
else:
|
|
# Not retryable or retry budget exhausted
|
|
if status_code in ctx.retry_codes:
|
|
logger.error(
|
|
f"Retry exhausted after {ctx.attempt} attempts, "
|
|
f"last status: {status_code}, total delay: {ctx.total_delay:.2f}s"
|
|
)
|
|
else:
|
|
logger.error(f"Non-retryable status code: {status_code}")
|
|
|
|
# Raise last failed exception
|
|
raise
|
|
|
|
|
|
__all__ = [
|
|
"RetryContext",
|
|
"retry_on_status",
|
|
"extract_retry_after",
|
|
]
|