crawler-intra-mart/utils/retry.py

import asyncio
from functools import wraps
import random

from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
    before_sleep_log
)

from utils.logger import logger


class RetryConfig:
    """
    Global retry configuration.
    """

    DEFAULT_ATTEMPTS = 3

    DEFAULT_MIN_WAIT = 1

    DEFAULT_MAX_WAIT = 10
    
    DEFAULT_JITTER = 0.5


def retry_sync(
    attempts=RetryConfig.DEFAULT_ATTEMPTS,
    min_wait=RetryConfig.DEFAULT_MIN_WAIT,
    max_wait=RetryConfig.DEFAULT_MAX_WAIT,
    exceptions=(Exception,)
):
    """
    Retry decorator for sync functions.

    Example:

    @retry_sync()
    def task():
        ...
    """

    return retry(
        stop=stop_after_attempt(attempts),

        wait=wait_exponential(
            multiplier=1,
            min=min_wait,
            max=max_wait
        ),

        retry=retry_if_exception_type(
            exceptions
        ),

        before_sleep=before_sleep_log(
            logger,
            "WARNING"
        ),

        reraise=True
    )


def retry_async(
    attempts=RetryConfig.DEFAULT_ATTEMPTS,
    min_wait=RetryConfig.DEFAULT_MIN_WAIT,
    max_wait=RetryConfig.DEFAULT_MAX_WAIT,
    jitter=RetryConfig.DEFAULT_JITTER,
    exceptions=(Exception,)
):
    """
    Retry decorator for async functions.

    Example:

    @retry_async()
    async def fetch():
        ...
    """

    def decorator(func):

        @wraps(func)
        async def wrapper(*args, **kwargs):

            last_exception = None

            for attempt in range(
                1,
                attempts + 1
            ):

                try:

                    logger.debug(
                        f"Attempt {attempt}/"
                        f"{attempts}: "
                        f"{func.__name__}"
                    )

                    return await func(
                        *args,
                        **kwargs
                    )

                except exceptions as e:

                    last_exception = e

                    logger.warning(
                        f"Retry failed "
                        f"{attempt}/{attempts} "
                        f"for {func.__name__}: {e}"
                    )

                    if attempt >= attempts:
                        break

                    base_sleep = min(
                        min_wait * (2 ** (attempt - 1)),
                        max_wait
                    )

                    sleep_time = base_sleep + random.uniform(
                        0,
                        jitter
                    )

                    logger.info(
                        f"Sleeping "
                    f"{sleep_time:.2f}s before retry"
                    )

                    await asyncio.sleep(
                        sleep_time
                    )

            logger.error(
                f"Max retries exceeded "
                f"for {func.__name__}"
            )

            raise last_exception

        return wrapper

    return decorator


async def retry_operation(
    coro,
    attempts=3,
    delay=1,
    exceptions=(Exception,)
):
    """
    Retry arbitrary async coroutine.

    Example:

    result = await retry_operation(
        fetch_page(url)
    )
    """

    last_exception = None

    for attempt in range(
        1,
        attempts + 1
    ):

        try:

            logger.debug(
                f"Retry operation "
                f"attempt {attempt}/{attempts}"
            )

            return await coro

        except exceptions as e:

            last_exception = e

            logger.warning(
                f"Retry operation failed "
                f"{attempt}/{attempts}: {e}"
            )

            if attempt >= attempts:
                break

            await asyncio.sleep(delay)

    logger.error(
        "Retry operation exceeded "
        "max attempts."
    )

    raise last_exception


def safe_execute(
    func,
    *args,
    default=None,
    log_error=True,
    **kwargs
):
    """
    Execute function safely.

    Returns default value if failed.

    Example:

    result = safe_execute(
        parse_html,
        html,
        default=""
    )
    """

    try:

        return func(
            *args,
            **kwargs
        )

    except Exception as e:

        if log_error:

            logger.exception(
                f"Safe execute failed "
                f"{func.__name__}: {e}"
            )

        return default


async def safe_execute_async(
    func,
    *args,
    default=None,
    log_error=True,
    **kwargs
):
    """
    Safe async execution.
    """

    try:

        return await func(
            *args,
            **kwargs
        )

    except Exception as e:

        if log_error:

            logger.exception(
                f"Safe async execute failed "
                f"{func.__name__}: {e}"
            )

        return default