TrueNAS Core MCP Server

"""Retry policies with exponential backoff."""

import asyncio
import logging
from typing import Callable, Optional, Tuple, Type
from functools import wraps
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class RetryPolicy:
    """Retry policy configuration."""

    max_attempts: int = 3
    initial_delay: float = 1.0
    max_delay: float = 60.0
    exponential_base: float = 2.0
    jitter: bool = True
    retryable_exceptions: Tuple[Type[Exception], ...] = (Exception,)


async def exponential_backoff(
    func: Callable,
    policy: Optional[RetryPolicy] = None,
    *args,
    **kwargs,
):
    """
    Execute function with exponential backoff retry.

    Args:
        func: Async function to execute
        policy: Retry policy configuration
        *args: Function arguments
        **kwargs: Function keyword arguments

    Returns:
        Function result

    Raises:
        Last exception if all retries exhausted
    """
    policy = policy or RetryPolicy()
    last_exception = None

    for attempt in range(policy.max_attempts):
        try:
            return await func(*args, **kwargs)
        except policy.retryable_exceptions as e:
            last_exception = e

            if attempt + 1 >= policy.max_attempts:
                logger.error(f"All {policy.max_attempts} retry attempts exhausted")
                raise

            # Calculate delay with exponential backoff
            delay = min(
                policy.initial_delay * (policy.exponential_base**attempt),
                policy.max_delay,
            )

            # Add jitter if enabled
            if policy.jitter:
                import random

                delay = delay * (0.5 + random.random())

            logger.warning(
                f"Attempt {attempt + 1}/{policy.max_attempts} failed, "
                f"retrying in {delay:.2f}s: {e}"
            )
            await asyncio.sleep(delay)

    # Should never reach here, but just in case
    if last_exception:
        raise last_exception


def retry(
    max_attempts: int = 3,
    initial_delay: float = 1.0,
    max_delay: float = 60.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
    retryable_exceptions: Tuple[Type[Exception], ...] = (Exception,),
):
    """
    Decorator for automatic retry with exponential backoff.

    Args:
        max_attempts: Maximum retry attempts
        initial_delay: Initial delay in seconds
        max_delay: Maximum delay in seconds
        exponential_base: Base for exponential backoff
        jitter: Whether to add random jitter
        retryable_exceptions: Tuple of exception types to retry on

    Example:
        @retry(max_attempts=5, initial_delay=0.5)
        async def flaky_api_call():
            ...
    """
    policy = RetryPolicy(
        max_attempts=max_attempts,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_base=exponential_base,
        jitter=jitter,
        retryable_exceptions=retryable_exceptions,
    )

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            return await exponential_backoff(func, policy, *args, **kwargs)

        return wrapper

    return decorator