Antoni Nuñez Romeu
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
from functools import wraps
|
|
from typing import Any, Callable, TypeVar
|
|
|
|
from src.core.exceptions import ConnectorException
|
|
from src.core.logging import get_logger
|
|
|
|
logger = get_logger("connector_retry")
|
|
|
|
T = TypeVar("T")
|
|
|
|
DEFAULT_RETRYABLE_EXCEPTIONS = (ConnectionError, TimeoutError, OSError)
|
|
|
|
|
|
@dataclass
|
|
class RetryPolicy:
|
|
max_attempts: int = 3
|
|
base_delay: float = 1.0
|
|
max_delay: float = 30.0
|
|
exponential_base: float = 2.0
|
|
jitter: bool = True
|
|
retryable_exceptions: tuple[type[Exception], ...] = field(default_factory=lambda: DEFAULT_RETRYABLE_EXCEPTIONS)
|
|
|
|
def delay_for_attempt(self, attempt: int) -> float:
|
|
import random
|
|
|
|
delay = min(self.base_delay * (self.exponential_base ** (attempt - 1)), self.max_delay)
|
|
if self.jitter:
|
|
delay = delay * (0.5 + random.random() * 0.5)
|
|
return delay
|
|
|
|
def should_retry(self, exception: Exception) -> bool:
|
|
if isinstance(exception, ConnectorException):
|
|
return False
|
|
return isinstance(exception, self.retryable_exceptions)
|
|
|
|
|
|
async def with_retry(
|
|
func: Callable[..., Any],
|
|
*args: Any,
|
|
policy: RetryPolicy | None = None,
|
|
connector_type: str = "unknown",
|
|
operation: str = "unknown",
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
p = policy or RetryPolicy()
|
|
last_exception: Exception | None = None
|
|
|
|
for attempt in range(1, p.max_attempts + 1):
|
|
try:
|
|
result = await func(*args, **kwargs)
|
|
if attempt > 1:
|
|
logger.info("retry_succeeded", extra={"connector_type": connector_type, "operation": operation, "attempt": attempt})
|
|
return result
|
|
except Exception as exc:
|
|
last_exception = exc
|
|
if not p.should_retry(exc) or attempt >= p.max_attempts:
|
|
break
|
|
delay = p.delay_for_attempt(attempt)
|
|
logger.warning(
|
|
"retry_attempt",
|
|
extra={
|
|
"connector_type": connector_type,
|
|
"operation": operation,
|
|
"attempt": attempt,
|
|
"delay_s": round(delay, 2),
|
|
"error": str(exc),
|
|
},
|
|
)
|
|
await asyncio.sleep(delay)
|
|
|
|
raise last_exception # type: ignore[misc]
|
|
|
|
|
|
def retry(policy: RetryPolicy | None = None):
|
|
p = policy or RetryPolicy()
|
|
|
|
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
@wraps(func)
|
|
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
connector_type = getattr(args[0], "connector_type", "unknown") if args else "unknown"
|
|
return await with_retry(func, *args, policy=p, connector_type=connector_type, operation=func.__name__, **kwargs)
|
|
return wrapper
|
|
return decorator
|