pyinfra/mini_queue/utils/meta.py

import functools
import logging
import time
from functools import partial, wraps
from math import exp
from typing import Tuple, Type, Callable


def invocation_counter(var: str):
    def inner(func):
        invocation_count = 0

        @functools.wraps(func)
        def inner(*args, **kwargs):
            nonlocal invocation_count
            invocation_count += 1

            if var not in kwargs:
                kwargs[var] = invocation_count

            return func(*args, **kwargs)

        return inner

    return inner


class NoAttemptsLeft(Exception):
    pass


class MaxTimeoutReached(Exception):
    pass


class _MethodDecoratorAdaptor(object):
    def __init__(self, decorator, func):
        self.decorator = decorator
        self.func = func

    def __call__(self, *args, **kwargs):
        return self.decorator(self.func)(*args, **kwargs)

    def __get__(self, obj, objtype):
        return partial(self.__call__, obj)


def auto_adapt_to_methods(decorator):
    """Allows you to use the same decorator on methods and functions,
    hiding the self argument from the decorator."""

    def adapt(func):
        return _MethodDecoratorAdaptor(decorator, func)

    return adapt


def max_attempts(
    n_attempts: int = 5, exceptions: Tuple[Type[Exception]] = None, timeout: float = 0.1, max_timeout: float = 10
) -> Callable:
    """Function decorator that attempts to run the wrapped function a certain number of times. Timeouts increase
    exponentially according to `Tₖ ≔ t eᵏ`, where `t` is the timeout factor `timeout` and `k` is the attempt number.
    If `∑ᵢ Tᵢ > mₜ` at the `i-th` attempt, where `mₜ` is the maximum timeout, then the function raises
    MaxTimeoutReached. If `k > mₐ`, where `mₐ` is the maximum number of attempts allowed, then the function
    raises NoAttemptsLeft.

    Args:
        n_attempts: Number of times to attempt running the wrapped function.
        exceptions: Exceptions to catch for a re-attempt.
        timeout: Timeout factor in seconds.
        max_timeout: Maximum allowed timeout.

    Raises:
        MaxTimeoutReached
        NoAttemptsLeft

    Returns:
        Wrapped function.
    """
    if not exceptions:
        exceptions = (Exception,)
    assert isinstance(exceptions, tuple)

    @auto_adapt_to_methods
    def inner(func):
        @wraps(func)
        def inner(*args, **kwargs):
            def run_attempt(attempt, timeout_aggr=0):
                if attempt:
                    try:
                        return func(*args, **kwargs)
                    except exceptions as err:

                        attempt_num = n_attempts - attempt + 1
                        next_timeout = timeout * exp(attempt_num - 1)  # start with timeout * e^0 = timeout

                        logging.warn(f"{func.__name__} failed; attempt {attempt_num} of {n_attempts}")

                        time_left = max(0, max_timeout - timeout_aggr)
                        if time_left:
                            sleep_for = min(next_timeout, time_left)
                            time.sleep(sleep_for)
                            return run_attempt(attempt - 1, timeout_aggr + sleep_for)
                        else:
                            logging.exception(err)
                            raise MaxTimeoutReached(
                                f"{func.__name__} reached maximum timeout ({max_timeout}) after {attempt_num} attempts."
                            )
                else:
                    raise NoAttemptsLeft(f"{func.__name__} failed {n_attempts} times; all attempts expended.")

            return run_attempt(n_attempts)

        return inner

    return inner