1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""
5Function/method decorators that provide timeout and retry logic.
6"""
7
8import functools
9import itertools
10import sys
11
12import six
13
14from devil.android import device_errors
15from devil.utils import cmd_helper
16from devil.utils import reraiser_thread
17from devil.utils import timeout_retry
18
19DEFAULT_TIMEOUT_ATTR = '_default_timeout'
20DEFAULT_RETRIES_ATTR = '_default_retries'
21
22
23def _TimeoutRetryWrapper(f,
24                         timeout_func,
25                         retries_func,
26                         retry_if_func=timeout_retry.AlwaysRetry,
27                         pass_values=False):
28  """ Wraps a funcion with timeout and retry handling logic.
29
30  Args:
31    f: The function to wrap.
32    timeout_func: A callable that returns the timeout value.
33    retries_func: A callable that returns the retries value.
34    pass_values: If True, passes the values returned by |timeout_func| and
35                 |retries_func| to the wrapped function as 'timeout' and
36                 'retries' kwargs, respectively.
37  Returns:
38    The wrapped function.
39  """
40
41  @functools.wraps(f)
42  def timeout_retry_wrapper(*args, **kwargs):
43    timeout = timeout_func(*args, **kwargs)
44    retries = retries_func(*args, **kwargs)
45    if pass_values:
46      kwargs['timeout'] = timeout
47      kwargs['retries'] = retries
48
49    @functools.wraps(f)
50    def impl():
51      return f(*args, **kwargs)
52
53    try:
54      if timeout_retry.CurrentTimeoutThreadGroup():
55        # Don't wrap if there's already an outer timeout thread.
56        return impl()
57      else:
58        desc = '%s(%s)' % (f.__name__, ', '.join(
59            itertools.chain(
60                (str(a) for a in args),
61                ('%s=%s' % (k, str(v)) for k, v in six.iteritems(kwargs)))))
62        return timeout_retry.Run(
63            impl, timeout, retries, desc=desc, retry_if_func=retry_if_func)
64    except reraiser_thread.TimeoutError as e:
65      six.reraise(
66          device_errors.CommandTimeoutError,
67          device_errors.CommandTimeoutError(str(e)),
68          sys.exc_info()[2])
69    except cmd_helper.TimeoutError as e:
70      six.reraise(
71          device_errors.CommandTimeoutError,
72          device_errors.CommandTimeoutError(str(e), output=e.output),
73          sys.exc_info()[2])
74
75  return timeout_retry_wrapper
76
77
78def WithTimeoutAndRetries(f):
79  """A decorator that handles timeouts and retries.
80
81  'timeout' and 'retries' kwargs must be passed to the function.
82
83  Args:
84    f: The function to decorate.
85  Returns:
86    The decorated function.
87  """
88  get_timeout = lambda *a, **kw: kw['timeout']
89  get_retries = lambda *a, **kw: kw['retries']
90  return _TimeoutRetryWrapper(f, get_timeout, get_retries)
91
92
93def WithTimeoutAndConditionalRetries(retry_if_func):
94  """Returns a decorator that handles timeouts and, in some cases, retries.
95
96  'timeout' and 'retries' kwargs must be passed to the function.
97
98  Args:
99    retry_if_func: A unary callable that takes an exception and returns
100      whether failures should be retried.
101  Returns:
102    The actual decorator.
103  """
104
105  def decorator(f):
106    get_timeout = lambda *a, **kw: kw['timeout']
107    get_retries = lambda *a, **kw: kw['retries']
108    return _TimeoutRetryWrapper(
109        f, get_timeout, get_retries, retry_if_func=retry_if_func)
110
111  return decorator
112
113
114def WithExplicitTimeoutAndRetries(timeout, retries):
115  """Returns a decorator that handles timeouts and retries.
116
117  The provided |timeout| and |retries| values are always used.
118
119  Args:
120    timeout: The number of seconds to wait for the decorated function to
121             return. Always used.
122    retries: The number of times the decorated function should be retried on
123             failure. Always used.
124  Returns:
125    The actual decorator.
126  """
127
128  def decorator(f):
129    get_timeout = lambda *a, **kw: timeout
130    get_retries = lambda *a, **kw: retries
131    return _TimeoutRetryWrapper(f, get_timeout, get_retries)
132
133  return decorator
134
135
136def WithTimeoutAndRetriesDefaults(default_timeout, default_retries):
137  """Returns a decorator that handles timeouts and retries.
138
139  The provided |default_timeout| and |default_retries| values are used only
140  if timeout and retries values are not provided.
141
142  Args:
143    default_timeout: The number of seconds to wait for the decorated function
144                     to return. Only used if a 'timeout' kwarg is not passed
145                     to the decorated function.
146    default_retries: The number of times the decorated function should be
147                     retried on failure. Only used if a 'retries' kwarg is not
148                     passed to the decorated function.
149  Returns:
150    The actual decorator.
151  """
152
153  def decorator(f):
154    get_timeout = lambda *a, **kw: kw.get('timeout', default_timeout)
155    get_retries = lambda *a, **kw: kw.get('retries', default_retries)
156    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)
157
158  return decorator
159
160
161def WithTimeoutAndRetriesFromInstance(default_timeout_name=DEFAULT_TIMEOUT_ATTR,
162                                      default_retries_name=DEFAULT_RETRIES_ATTR,
163                                      min_default_timeout=None):
164  """Returns a decorator that handles timeouts and retries.
165
166  The provided |default_timeout_name| and |default_retries_name| are used to
167  get the default timeout value and the default retries value from the object
168  instance if timeout and retries values are not provided.
169
170  Note that this should only be used to decorate methods, not functions.
171
172  Args:
173    default_timeout_name: The name of the default timeout attribute of the
174                          instance.
175    default_retries_name: The name of the default retries attribute of the
176                          instance.
177    min_timeout: Miniumum timeout to be used when using instance timeout.
178  Returns:
179    The actual decorator.
180  """
181
182  def decorator(f):
183    def get_timeout(inst, *_args, **kwargs):
184      ret = getattr(inst, default_timeout_name)
185      if min_default_timeout is not None:
186        ret = max(min_default_timeout, ret)
187      return kwargs.get('timeout', ret)
188
189    def get_retries(inst, *_args, **kwargs):
190      return kwargs.get('retries', getattr(inst, default_retries_name))
191
192    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)
193
194  return decorator
195