Skip to main content
`Ctrl`+`K`

# Source code for silx.utils.retry

    # /*##########################################################################
    # Copyright (C) 2016-2023 European Synchrotron Radiation Facility
    #
    # Permission is hereby granted, free of charge, to any person obtaining a copy
    # of this software and associated documentation files (the "Software"), to deal
    # in the Software without restriction, including without limitation the rights
    # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    # copies of the Software, and to permit persons to whom the Software is
    # furnished to do so, subject to the following conditions:
    #
    # The above copyright notice and this permission notice shall be included in
    # all copies or substantial portions of the Software.
    #
    # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    # THE SOFTWARE.
    #
    # ############################################################################*/
    """
    This module provides utility methods for retrying methods until they
    no longer fail.
    """

    __authors__ = ["W. de Nolf"]
    __license__ = "MIT"
    __date__ = "28/11/2023"

    import time
    import inspect
    from functools import wraps
    from contextlib import contextmanager
    from queue import Empty

    RETRY_PERIOD = 0.01

    class RetryTimeoutError(TimeoutError):
        pass

    class RetryError(RuntimeError):
        pass

    def _default_retry_on_error(e):
        """
        :param BaseException e:
        :returns bool:
        """
        return isinstance(e, RetryError)

    @contextmanager
    def _handle_exception(retry_state):
        try:
            yield
        except BaseException as e:
            retry_on_error = retry_state.get("retry_on_error")
            if retry_on_error is not None and retry_on_error(e):
                retry_state["exception"] = e
            else:
                raise

    def _retry_loop(retry_timeout=None, retry_period=None, retry_on_error=None):
        """Iterator which is endless or ends with an RetryTimeoutError.
        It yields a dictionary which can be used to influence the loop.

        :param num retry_timeout:
        :param num retry_period: sleep before retry
        :param callable or None retry_on_error: checks whether an exception is
                                                eligible for retry
        """
        has_timeout = retry_timeout is not None
        if has_timeout:
            t0 = time.time()
        else:
            t0 = None
        retry_state = {"t0": t0, "exception": None, "retry_on_error": retry_on_error}
        while True:
            yield retry_state
            if has_timeout and (time.time() - retry_state["t0"]) > retry_timeout:
                err_msg = "%s seconds" % retry_timeout
                cause = retry_state.get("exception")
                raise RetryTimeoutError(err_msg) from cause
            if retry_period is not None:
                time.sleep(retry_period)

    def retry(
        retry_timeout=None, retry_period=None, retry_on_error=_default_retry_on_error
    ):
        """Decorator for a method that needs to be executed until it not longer
        fails or until `retry_on_error` returns False.

        The decorator arguments can be overriden by using them when calling the
        decorated method.

        Generator functions are required to have a `start_index` argument which allows
        the method to start iterating from the last failure when called on retry.

        :param num retry_timeout:
        :param num retry_period: sleep before retry
        :param callable or None retry_on_error: checks whether an exception is
                                                eligible for retry
        """

        if retry_period is None:
            retry_period = RETRY_PERIOD

        def decorator(method):
            if inspect.isgeneratorfunction(method):
                if "start_index" not in inspect.signature(method).parameters:
                    raise TypeError(
                        "The generator function '%s' needs a `start_index` named argument because it is wrapped with the `retry` decorator."
                        % method.__name__
                    )

                @wraps(method)
                def wrapper(*args, **kw):
                    _retry_timeout = kw.pop("retry_timeout", retry_timeout)
                    _retry_period = kw.pop("retry_period", retry_period)
                    _retry_on_error = kw.pop("retry_on_error", retry_on_error)
                    start_index = kw.pop("start_index", 0)
                    if start_index is None:
                        start_index = 0
                    for retry_state in _retry_loop(
                        retry_timeout=_retry_timeout,
                        retry_period=_retry_period,
                        retry_on_error=_retry_on_error,
                    ):
                        with _handle_exception(retry_state):
                            oretry_on_error = retry_state["retry_on_error"]
                            for result in method(*args, start_index=start_index, **kw):
                                start_index += 1
                                retry_state["retry_on_error"] = None
                                # any exception here will NOT cause a retry
                                yield result
                                # restart the retry loop
                                if retry_state["t0"] is not None:
                                    retry_state["t0"] = time.time()
                                retry_state["retry_on_error"] = oretry_on_error
                            return

            else:

                @wraps(method)
                def wrapper(*args, **kw):
                    _retry_timeout = kw.pop("retry_timeout", retry_timeout)
                    _retry_period = kw.pop("retry_period", retry_period)
                    _retry_on_error = kw.pop("retry_on_error", retry_on_error)
                    for retry_state in _retry_loop(
                        retry_timeout=_retry_timeout,
                        retry_period=_retry_period,
                        retry_on_error=_retry_on_error,
                    ):
                        with _handle_exception(retry_state):
                            return method(*args, **kw)

            return wrapper

        return decorator

    def retry_contextmanager(
        retry_timeout=None, retry_period=None, retry_on_error=_default_retry_on_error
    ):
        """Decorator to make a context manager from a method that needs to be
        entered until it no longer fails or until `retry_on_error` returns False.

        The decorator arguments can be overriden by using them when calling the
        decorated method.

        :param num retry_timeout:
        :param num retry_period: sleep before retry
        :param callable or None retry_on_error: checks whether an exception is
                                                eligible for retry
        """

        if retry_period is None:
            retry_period = RETRY_PERIOD

        def decorator(method):
            @wraps(method)
            def wrapper(*args, **kw):
                _retry_timeout = kw.pop("retry_timeout", retry_timeout)
                _retry_period = kw.pop("retry_period", retry_period)
                _retry_on_error = kw.pop("retry_on_error", retry_on_error)
                for retry_state in _retry_loop(
                    retry_timeout=_retry_timeout,
                    retry_period=_retry_period,
                    retry_on_error=_retry_on_error,
                ):
                    with _handle_exception(retry_state):
                        ctx = method(*args, **kw)
                        result = next(ctx)
                        retry_state["retry_on_error"] = None
                        yield result
                        try:
                            next(ctx)
                        except StopIteration:
                            return
                        else:
                            raise RuntimeError(str(method) + " should only yield once")

            return contextmanager(wrapper)

        return decorator

    def _subprocess_main(queue, method, retry_on_error, *args, **kw):
        try:
            result = method(*args, **kw)
        except BaseException as e:
            if retry_on_error(e):
                # As the traceback gets lost, make sure the top-level
                # exception is RetryError
                e = RetryError(str(e))
            queue.put(e)
        else:
            queue.put(result)

    def retry_in_subprocess(
        retry_timeout=None, retry_period=None, retry_on_error=_default_retry_on_error
    ):
        """Same as `retry` but it also retries segmentation faults.

        As subprocesses are spawned, you cannot use this decorator with the "@" syntax
        because the decorated method needs to be an attribute of a module:

        .. code-block:: python

            def _method(*args, **kw):
                ...

            method = retry_in_subprocess()(_method)

        :param num retry_timeout:
        :param num retry_period: sleep before retry
        :param callable or None retry_on_error: checks whether an exception is
                                                eligible for retry
        """

        if retry_period is None:
            retry_period = RETRY_PERIOD

        import multiprocessing

        def decorator(method):
            @wraps(method)
            def wrapper(*args, **kw):
                _retry_timeout = kw.pop("retry_timeout", retry_timeout)
                _retry_period = kw.pop("retry_period", retry_period)
                _retry_on_error = kw.pop("retry_on_error", retry_on_error)

                ctx = multiprocessing.get_context("spawn")

                def start_subprocess():
                    queue = ctx.Queue(maxsize=1)
                    p = ctx.Process(
                        target=_subprocess_main,
                        args=(queue, method, retry_on_error) + args,
                        kwargs=kw,
                    )
                    p.start()
                    return p, queue

                def stop_subprocess(p):
                    try:
                        p.kill()
                    except AttributeError:
                        p.terminate()
                    p.join()

                p, queue = start_subprocess()
                try:
                    for retry_state in _retry_loop(
                        retry_timeout=_retry_timeout, retry_on_error=_retry_on_error
                    ):
                        with _handle_exception(retry_state):
                            if not p.is_alive():
                                p, queue = start_subprocess()
                            try:
                                result = queue.get(block=True, timeout=_retry_period)
                            except Empty:
                                pass
                            except ValueError:
                                pass
                            else:
                                if isinstance(result, BaseException):
                                    stop_subprocess(p)
                                    raise result
                                else:
                                    return result
                finally:
                    stop_subprocess(p)

            return wrapper

        return decorator