1#
2# Copyright (C) 2016 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16import sys
17import time
18
19def GenericRetry(handler, max_retry, functor,
20                 sleep=0, backoff_factor=1, success_functor=lambda x: None,
21                 raise_first_exception_on_failure=True, *args, **kwargs):
22    """Generic retry loop w/ optional break out depending on exceptions.
23
24    To retry based on the return value of |functor| see the timeout_util module.
25
26    Keep in mind that the total sleep time will be the triangular value of
27    max_retry multiplied by the sleep value.  e.g. max_retry=5 and sleep=10
28    will be T5 (i.e. 5+4+3+2+1) times 10, or 150 seconds total.  Rather than
29    use a large sleep value, you should lean more towards large retries and
30    lower sleep intervals, or by utilizing backoff_factor.
31
32    Args:
33        handler: A functor invoked w/ the exception instance that
34            functor(*args, **kwargs) threw.  If it returns True, then a
35            retry is attempted.  If False, the exception is re-raised.
36        max_retry: A positive integer representing how many times to retry
37            the command before giving up.  Worst case, the command is invoked
38            max_retry + 1) times before failing.
39        functor: A callable to pass args and kwargs to.
40        sleep: Optional keyword.  Multiplier for how long to sleep between
41            retries; will delay (1*sleep) the first time, then (2*sleep),
42            continuing via attempt * sleep.
43        backoff_factor: Optional keyword. If supplied and > 1, subsequent sleeps
44                        will be of length (backoff_factor ^ (attempt - 1)) * sleep,
45                        rather than the default behavior of attempt * sleep.
46        success_functor: Optional functor that accepts 1 argument. Will be called
47                         after successful call to |functor|, with the argument
48                         being the number of attempts (1 = |functor| succeeded on
49                         first try).
50        raise_first_exception_on_failure: Optional boolean which determines which
51                                          exception is raised upon failure after
52                                          retries. If True, the first exception
53                                          that was encountered. If False, the
54                                          final one. Default: True.
55        *args: Positional args passed to functor.
56        **kwargs: Optional args passed to functor.
57
58    Returns:
59        Whatever functor(*args, **kwargs) returns.
60
61    Raises:
62        Exception: Whatever exceptions functor(*args, **kwargs) throws and
63            isn't suppressed is raised.  Note that the first exception encountered
64            is what's thrown.
65    """
66
67    if max_retry < 0:
68        raise ValueError('max_retry needs to be zero or more: %s' % max_retry)
69
70    if backoff_factor < 1:
71        raise ValueError('backoff_factor must be 1 or greater: %s'
72                          % backoff_factor)
73
74    ret, success = (None, False)
75    attempt = 0
76
77    exc_info = None
78    for attempt in xrange(max_retry + 1):
79        if attempt and sleep:
80            if backoff_factor > 1:
81                sleep_time = sleep * backoff_factor ** (attempt - 1)
82            else:
83                sleep_time = sleep * attempt
84            time.sleep(sleep_time)
85        try:
86            ret = functor(*args, **kwargs)
87            success = True
88            break
89        except Exception as e:
90            # Note we're not snagging BaseException, so MemoryError/KeyboardInterrupt
91            # and friends don't enter this except block.
92            if not handler(e):
93                raise
94            # If raise_first_exception_on_failure, we intentionally ignore
95            # any failures in later attempts since we'll throw the original
96            # failure if all retries fail.
97            if exc_info is None or not raise_first_exception_on_failure:
98                exc_info = sys.exc_info()
99
100    if success:
101        success_functor(attempt + 1)
102        return ret
103
104    raise exc_info[0], exc_info[1], exc_info[2]
105
106def RetryException(exc_retry, max_retry, functor, *args, **kwargs):
107    """Convenience wrapper for GenericRetry based on exceptions.
108
109    Args:
110        exc_retry: A class (or tuple of classes).  If the raised exception
111            is the given class(es), a retry will be attempted.  Otherwise,
112            the exception is raised.
113        max_retry: See GenericRetry.
114        functor: See GenericRetry.
115        *args: See GenericRetry.
116        **kwargs: See GenericRetry.
117
118    Returns:
119        Return what functor returns.
120
121    Raises:
122        TypeError, if exc_retry is of an unexpected type.
123    """
124    if not isinstance(exc_retry, (tuple, type)):
125        raise TypeError("exc_retry should be an exception (or tuple), not %r" %
126                         exc_retry)
127    def _Handler(exc, values=exc_retry):
128        return isinstance(exc, values)
129    return GenericRetry(_Handler, max_retry, functor, *args, **kwargs)
130