Coverage for src/taipanstack/resilience/retry.py: 100%
229 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 15:01 +0000
1"""
2Retry logic with exponential backoff.
4Provides decorators for automatic retry of failing operations
5with configurable backoff strategies. Compatible with any
6Python framework (sync and async).
7"""
9import asyncio
10import functools
11import inspect
12import logging
13import math
14import secrets
15import time
16from collections.abc import Awaitable, Callable
17from dataclasses import dataclass
18from types import TracebackType
19from typing import NoReturn, ParamSpec, Protocol, TypeVar, cast, overload
21from taipanstack.core.result import Err
23P = ParamSpec("P")
24R = TypeVar("R")
27def _handle_retry_exception(
28 e: Exception,
29 attempt: int,
30 func_name: str,
31 valid_on: tuple[type[Exception], ...] | type[Exception],
32 config: "RetryConfig",
33) -> tuple[bool, float]:
34 """Handle a retry attempt exception.
36 Args:
37 e: The exception raised.
38 attempt: The current attempt number.
39 func_name: The name of the function.
40 valid_on: The exception types to retry on.
41 config: The retry configuration.
43 Returns:
44 True if the exception should be retried, False otherwise.
46 Raises:
47 Exception: If the exception is not matched and should not be retried.
49 """
50 try:
51 is_match = isinstance(e, valid_on)
52 except TypeError:
53 is_match = False
54 if not is_match:
55 raise e
57 if attempt == config.max_attempts:
58 _log_all_failed(func_name, e, config)
59 return False, 0.0
61 delay = calculate_delay(attempt, config)
62 _log_retry_attempt(func_name, attempt, e, delay, config)
63 return True, delay
66class RetryDecorator(Protocol):
67 """Protocol for the retry decorator."""
69 @overload
70 def __call__(self, func: Callable[P, R]) -> Callable[P, R]: ...
72 @overload
73 def __call__(
74 self,
75 func: Callable[P, Awaitable[R]],
76 ) -> Callable[P, Awaitable[R]]: ...
79logger = logging.getLogger("taipanstack.resilience.retry")
81try:
82 import structlog as _structlog
84 _structlog_logger = _structlog.get_logger("taipanstack.resilience.retry")
85 _HAS_STRUCTLOG = True
86except ImportError:
87 _structlog_logger = None
88 _HAS_STRUCTLOG = False
91def _validate_finite_or_default(
92 obj: object,
93 attr_name: str,
94 default_val: float | int,
95) -> None:
96 """Validate that an attribute is finite, falling back to a default."""
97 try:
98 val = cast(float | int, getattr(obj, attr_name))
99 if not math.isfinite(val):
100 raise ValueError(f"{attr_name} must be a finite number")
101 except TypeError:
102 object.__setattr__(obj, attr_name, default_val)
105@dataclass(frozen=True)
106class RetryConfig:
107 """Configuration for retry behavior.
109 Attributes:
110 max_attempts: Maximum number of retry attempts.
111 initial_delay: Initial delay between retries in seconds.
112 max_delay: Maximum delay between retries.
113 exponential_base: Base for exponential backoff (2 = double each time).
114 jitter: Whether to add random jitter to delays.
115 jitter_factor: Maximum jitter as fraction of delay (0.1 = 10%).
116 log_retries: Whether to emit standard log messages.
117 on_retry: Optional callback invoked on each retry.
119 """
121 max_attempts: int = 3
122 initial_delay: float = 1.0
123 max_delay: float = 60.0
124 exponential_base: float = 2.0
125 jitter: bool = True
126 jitter_factor: float = 0.1
127 log_retries: bool = True
128 on_retry: Callable[[int, int, Exception, float], None] | None = None
130 def __post_init__(self) -> None:
131 """Validate configuration parameters."""
132 _validate_finite_or_default(self, "max_attempts", 3)
133 _validate_finite_or_default(self, "initial_delay", 1.0)
134 _validate_finite_or_default(self, "max_delay", 60.0)
135 _validate_finite_or_default(self, "exponential_base", 2.0)
136 _validate_finite_or_default(self, "jitter_factor", 0.1)
139class RetryError(Exception):
140 """Raised when all retry attempts have failed."""
142 def __init__(
143 self,
144 message: str,
145 attempts: int,
146 last_exception: Exception | None = None,
147 ) -> None:
148 """Initialize RetryError.
150 Args:
151 message: Description of the retry failure.
152 attempts: Number of attempts made.
153 last_exception: The last exception that was raised.
155 """
156 self.attempts = attempts
157 self.last_exception = last_exception
158 super().__init__(message)
161def _calculate_base_delay(attempt: int, config: RetryConfig) -> float:
162 """Calculate base delay with exponential backoff."""
163 safe_attempt = max(1, attempt)
164 try:
165 delay = config.initial_delay * (config.exponential_base ** (safe_attempt - 1))
166 if not math.isfinite(delay):
167 delay = config.max_delay
168 except (OverflowError, TypeError):
169 delay = config.max_delay
171 try:
172 if not math.isfinite(delay):
173 delay = 0.0
174 return min(delay, config.max_delay)
175 except TypeError:
176 return 0.0
179def _compute_jitter_amount(delay: float, factor: float) -> float | None:
180 """Compute jitter amount safely."""
181 try:
182 amount = delay * factor
183 if math.isfinite(amount):
184 return amount
185 except (TypeError, OverflowError, ValueError, Exception) as e:
186 logger.warning("Failed to add jitter to delay due to mutation: %s", str(e))
187 return None
190def _apply_jitter(delay: float, config: RetryConfig) -> float:
191 """Apply jitter to delay."""
192 if not config.jitter or not math.isfinite(delay):
193 return delay
195 jitter_amount = _compute_jitter_amount(delay, config.jitter_factor)
196 if jitter_amount is not None:
197 try:
198 delay += secrets.SystemRandom().uniform(-jitter_amount, jitter_amount)
199 except Exception as e:
200 logger.warning("Failed to add jitter to delay: %s", str(e))
202 return delay
205def calculate_delay(
206 attempt: int,
207 config: RetryConfig,
208) -> float:
209 """Calculate delay before next retry.
211 Args:
212 attempt: Current attempt number (1-indexed).
213 config: Retry configuration.
215 Returns:
216 Delay in seconds before next retry.
218 """
219 delay = _calculate_base_delay(attempt, config)
220 delay = _apply_jitter(delay, config)
222 if not math.isfinite(delay) or delay < 0:
223 return 0.0
225 return delay
228def _log_retry_callback_failure(func_name: str, e: Exception) -> None:
229 """Log a failure during the retry callback execution."""
230 if _HAS_STRUCTLOG and _structlog_logger is not None:
231 _structlog_logger.error(
232 "retry_callback_failed",
233 function=func_name,
234 error=str(e),
235 )
236 else:
237 logger.error(
238 "Retry callback failed for %s: %s",
239 func_name,
240 str(e),
241 )
244def _log_retry_attempt_fallback(
245 func_name: str,
246 attempt: int,
247 exc: Exception,
248 delay: float,
249 config: RetryConfig,
250) -> None:
251 """Log the retry attempt if no callback is provided."""
252 if _HAS_STRUCTLOG and _structlog_logger is not None:
253 _structlog_logger.warning(
254 "retry_attempted",
255 function=func_name,
256 attempt=attempt,
257 max_attempts=config.max_attempts,
258 error=str(exc),
259 delay_seconds=round(delay, 3),
260 )
263def _invoke_retry_callback(
264 func_name: str,
265 attempt: int,
266 exc: Exception,
267 delay: float,
268 config: RetryConfig,
269) -> None:
270 """Invoke the retry callback if set, or emit structured log.
272 Args:
273 func_name: Name of the retried function.
274 attempt: Current attempt number.
275 exc: The exception that triggered the retry.
276 delay: Delay in seconds before the next attempt.
277 config: Retry configuration.
279 """
280 if config.on_retry is not None:
281 try:
282 config.on_retry(attempt, config.max_attempts, exc, delay)
283 except Exception as e:
284 _log_retry_callback_failure(func_name, e)
285 else:
286 _log_retry_attempt_fallback(func_name, attempt, exc, delay, config)
289def _log_retry_attempt(
290 func_name: str,
291 attempt: int,
292 exc: Exception,
293 delay: float,
294 config: RetryConfig,
295) -> None:
296 """Log a retry attempt via callback, structlog, or stdlib logger.
298 Args:
299 func_name: Name of the retried function.
300 attempt: Current attempt number.
301 exc: The exception that triggered the retry.
302 delay: Delay in seconds before the next attempt.
303 config: Retry configuration.
305 """
306 if config.log_retries:
307 logger.info(
308 "Attempt %d/%d failed for %s: %s. Retrying in %.2f seconds...",
309 attempt,
310 config.max_attempts,
311 func_name,
312 str(exc),
313 delay,
314 )
316 _invoke_retry_callback(func_name, attempt, exc, delay, config)
319def _log_all_failed(
320 func_name: str,
321 exc: Exception,
322 config: RetryConfig,
323) -> None:
324 """Log when all retry attempts have been exhausted.
326 Args:
327 func_name: Name of the retried function.
328 exc: The last exception raised.
329 config: Retry configuration.
331 """
332 if config.log_retries:
333 logger.warning(
334 "All %d attempts failed for %s: %s",
335 config.max_attempts,
336 func_name,
337 str(exc),
338 )
341def _raise_retry_error(
342 func_name: str,
343 max_attempts: int,
344 reraise: bool,
345 last_exception: Exception | None,
346) -> NoReturn:
347 """Raise a RetryError after all attempts fail.
349 Args:
350 func_name: Name of the retried function.
351 max_attempts: Number of attempts made.
352 reraise: Whether to reraise the original exception.
353 last_exception: The last exception that was raised.
355 Raises:
356 RetryError: The wrapped or unwrapped exception.
358 """
359 if reraise and last_exception is not None:
360 raise RetryError(
361 f"All {max_attempts} attempts failed for {func_name}",
362 attempts=max_attempts,
363 last_exception=last_exception,
364 ) from last_exception
366 raise RetryError(
367 f"All {max_attempts} attempts failed for {func_name}",
368 attempts=max_attempts,
369 last_exception=last_exception,
370 )
373def _ensure_tuple(
374 on: tuple[type[Exception], ...] | type[Exception],
375) -> tuple[type[Exception], ...]:
376 """Ensure the exception parameter is a tuple of exception types."""
377 if isinstance(on, type) and issubclass(on, BaseException):
378 return (on,)
379 if not isinstance(on, tuple):
380 msg = ( # type: ignore[unreachable]
381 "'on' parameter must be an exception class or a tuple of exception classes"
382 )
383 raise TypeError(msg)
384 return on
387def _validate_retry_exceptions(
388 on: tuple[type[Exception], ...] | type[Exception],
389) -> tuple[type[Exception], ...]:
390 """Validate that the 'on' parameter contains valid exception types."""
391 on_tuple = _ensure_tuple(on)
393 for exc_type in on_tuple:
394 if not isinstance(exc_type, type) or not issubclass(exc_type, BaseException):
395 msg = ( # type: ignore[unreachable]
396 f"All elements in 'on' must be subclasses of BaseException, "
397 f"got {type(exc_type).__name__}"
398 )
399 raise TypeError(msg)
401 return on_tuple
404def retry(
405 *,
406 max_attempts: int = 3,
407 initial_delay: float = 1.0,
408 max_delay: float = 60.0,
409 exponential_base: float = 2.0,
410 jitter: bool = True,
411 on: tuple[type[Exception], ...] | type[Exception] = (Exception,),
412 reraise: bool = True,
413 log_retries: bool = True,
414 on_retry: Callable[[int, int, Exception, float], None] | None = None,
415) -> RetryDecorator:
416 """Retry a sync or async function with exponential backoff.
418 Automatically retries the decorated function when specified
419 exceptions are raised, with configurable backoff strategy.
420 Detects coroutine functions and preserves their async nature.
422 Args:
423 max_attempts: Maximum number of retry attempts.
424 initial_delay: Initial delay between retries in seconds.
425 max_delay: Maximum delay between retries.
426 exponential_base: Base for exponential backoff.
427 jitter: Whether to add random jitter to delays.
428 on: Exception types to retry on.
429 reraise: Whether to reraise the last exception on failure.
430 log_retries: Whether to log retry attempts.
431 on_retry: Optional callback invoked on each retry with
432 (attempt, max_attempts, exception, delay). Useful for
433 custom monitoring or metrics collection.
435 Returns:
436 Decorated function with retry logic.
438 Example:
439 >>> @retry(max_attempts=3, on=(ConnectionError, TimeoutError))
440 ... def fetch_data(url: str) -> dict:
441 ... return requests.get(url, timeout=10).json()
443 >>> @retry(max_attempts=3, on_retry=lambda a, m, e, d: print(f"Retry {a}/{m}"))
444 ... def fragile_operation() -> str:
445 ... return do_something()
447 """
448 # Validate 'on' parameter at definition time (Fail-Fast)
449 valid_on = _validate_retry_exceptions(on)
451 config = RetryConfig(
452 max_attempts=max_attempts,
453 initial_delay=initial_delay,
454 max_delay=max_delay,
455 exponential_base=exponential_base,
456 jitter=jitter,
457 log_retries=log_retries,
458 on_retry=on_retry,
459 )
461 def decorator(
462 func: Callable[P, R] | Callable[P, Awaitable[R]],
463 ) -> Callable[P, R] | Callable[P, Awaitable[R]]:
464 if inspect.iscoroutinefunction(func):
465 func_coro = cast(Callable[P, Awaitable[R]], func)
467 @functools.wraps(func_coro)
468 async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
469 last_exception: Exception | None = None
470 last_result: R | None = None
472 for attempt in range(1, config.max_attempts + 1):
473 last_result = None
474 try:
475 last_result = await func_coro(*args, **kwargs)
476 if isinstance(last_result, Err):
477 err_val = last_result.unwrap_err()
478 if isinstance(err_val, valid_on):
479 raise err_val # noqa: TRY301
480 return last_result
481 except Exception as e:
482 last_exception = e
483 should_retry, delay = _handle_retry_exception(
484 e,
485 attempt,
486 cast(str, getattr(func_coro, "__name__", "unknown")),
487 valid_on,
488 config,
489 )
490 if not should_retry:
491 break
492 await asyncio.sleep(min(delay, 3600.0))
494 if last_result is not None and isinstance(last_result, Err):
495 return cast(R, last_result)
496 _raise_retry_error(
497 cast(str, getattr(func_coro, "__name__", "unknown")),
498 config.max_attempts,
499 reraise,
500 last_exception,
501 )
503 return async_wrapper
505 func_sync = cast(Callable[P, R], func)
507 @functools.wraps(func_sync)
508 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
509 last_exception: Exception | None = None
510 last_result: R | None = None
512 for attempt in range(1, config.max_attempts + 1):
513 last_result = None
514 try:
515 last_result = func_sync(*args, **kwargs)
516 if isinstance(last_result, Err):
517 err_val = last_result.unwrap_err()
518 if isinstance(err_val, valid_on):
519 raise err_val # noqa: TRY301
520 return last_result
521 except Exception as e:
522 last_exception = e
523 should_retry, delay = _handle_retry_exception(
524 e,
525 attempt,
526 cast(str, getattr(func_sync, "__name__", "unknown")),
527 valid_on,
528 config,
529 )
530 if not should_retry:
531 break
532 time.sleep(min(delay, 3600.0))
534 if last_result is not None and isinstance(last_result, Err):
535 return cast(R, last_result)
536 _raise_retry_error(
537 cast(str, getattr(func_sync, "__name__", "unknown")),
538 config.max_attempts,
539 reraise,
540 last_exception,
541 )
543 return wrapper
545 return cast(RetryDecorator, decorator)
548def retry_on_exception(
549 exception_types: tuple[type[Exception], ...],
550 max_attempts: int = 3,
551) -> RetryDecorator:
552 """Retry on specific exceptions.
554 A simpler alternative to the full retry decorator when you
555 just need basic retry functionality.
557 Args:
558 exception_types: Exception types to retry on.
559 max_attempts: Maximum number of attempts.
561 Returns:
562 Decorated function with retry logic.
564 Example:
565 >>> @retry_on_exception((ValueError,), max_attempts=2)
566 ... def parse_data(data: str) -> dict:
567 ... return json.loads(data)
569 """
570 return retry(
571 max_attempts=max_attempts,
572 on=exception_types,
573 jitter=False,
574 log_retries=False,
575 )
578class Retrier:
579 """Context manager for retry logic.
581 Provides a context manager interface for retry logic when
582 decorators are not suitable.
584 Example:
585 >>> retrier = Retrier(max_attempts=3, on=(ConnectionError,))
586 >>> with retrier:
587 ... result = some_operation()
589 """
591 def __init__(
592 self,
593 *,
594 max_attempts: int = 3,
595 initial_delay: float = 1.0,
596 max_delay: float = 60.0,
597 on: tuple[type[Exception], ...] = (Exception,),
598 ) -> None:
599 """Initialize Retrier.
601 Args:
602 max_attempts: Maximum retry attempts.
603 initial_delay: Initial delay between retries.
604 max_delay: Maximum delay between retries.
605 on: Exception types to retry on.
607 """
608 self.config = RetryConfig(
609 max_attempts=max_attempts,
610 initial_delay=initial_delay,
611 max_delay=max_delay,
612 )
613 self.exception_types = on
614 self.attempt = 0
615 self.last_exception: Exception | None = None
617 def __enter__(self) -> "Retrier":
618 """Enter the retry context."""
619 return self
621 def _increment_attempt(self) -> bool:
622 """Increment attempt counter safely."""
623 if (
624 not isinstance(self.attempt, (int, float))
625 or not math.isfinite(self.attempt)
626 or self.attempt < 0
627 ):
628 return False
630 self.attempt += 1
631 return True
633 def _should_retry(self, exc_type: type[BaseException] | None) -> bool:
634 """Determine if an exception should trigger a retry."""
635 if exc_type is None:
636 return False
637 try:
638 if not issubclass(exc_type, self.exception_types):
639 return False
640 except TypeError:
641 return False
643 if not self._increment_attempt():
644 return False
646 return self.attempt < self.config.max_attempts
648 def __exit__(
649 self,
650 exc_type: type[BaseException] | None,
651 exc_val: BaseException | None,
652 _exc_tb: TracebackType | None,
653 ) -> bool:
654 """Exit the retry context.
656 Returns True to suppress the exception if we should retry,
657 False to let it propagate.
658 """
659 # Safe cast: check inside _should_retry ensures we handle it right
660 if exc_val is not None:
661 self.last_exception = exc_val if isinstance(exc_val, Exception) else None
663 if not self._should_retry(exc_type):
664 return False
666 # Calculate delay and wait
667 delay = calculate_delay(self.attempt, self.config)
668 time.sleep(min(delay, 3600.0))
670 return True # Suppress exception and retry